from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import as_completed
from functools import reduce

from app.common.utils import task, init_executor
from app.globals import env
from app.macro import QueryDataCategory, TradeCode

executor = ThreadPoolExecutor(max_workers=10)
init_executor(executor)


@task(executor)
def query(account, category):
    worker = env.workers.get(account)
    if category == QueryDataCategory.BALANCE:
        return [worker.get_balance()]
    if category == QueryDataCategory.SHARES:
        return worker.get_shares()
    elif category == QueryDataCategory.TODAY_ENTRUSTMENT:
        return worker.get_entrustments()
    elif category == QueryDataCategory.NEW_STOCK:
        return worker.get_new_stocks()
    else:
        raise ValueError(f'not support category {category}')


@task(executor)
def buy_many(account, security, price, quantity_list):
    worker = env.workers.get(account)
    return worker.order(TradeCode.BUY, security, price, quantity_list)


@task(executor)
def sell_many(account, security, price, quantity_list):
    worker = env.workers.get(account)
    return worker.order(TradeCode.SELL, security, price, quantity_list)


@task(executor)
def cancel_many(account, security, hth_list):
    worker = env.workers.get(account)
    return worker.cancel(security, hth_list)


def batch_buy(security, price, lines):
    result = []
    for e in lines:
        account = e['account']
        quantity = e['quantity']
        result.append(buy_many.submit(account, security, price, quantity))
    ret = [r.result() for r in as_completed(result)]
    return reduce(lambda x, y: x + y, ret)


def batch_sell(security, price, lines):
    result = []
    for e in lines:
        account = e['account']
        quantity = e['quantity']
        result.append(sell_many.submit(account, security, price, quantity))
    ret = [r.result() for r in as_completed(result)]
    return reduce(lambda x, y: x + y, ret)


def batch_cancel(security, lines):
    result = []
    for e in lines:
        account = e['account']
        hth = e['hth']
        result.append(cancel_many.submit(account, security, hth))
    ret = [r.result() for r in result]
    return reduce(lambda x, y: x + y, ret)


def batch_query(category, accounts):
    result = []
    for account in accounts:
        if account not in env.workers:
            continue
        result.append(query.submit(account, category))
    ret = [r.result() for r in as_completed(result)]
    return ret and reduce(lambda x, y: x + y, ret) or []
